# import
import numpy as np
import pandas as pd

import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from util.models import *
from util.misc import *
from restrictiveness import *

## mat file
from scipy.io import loadmat

# load dataset
base = os.path.dirname(os.path.abspath(__file__))
file_name = os.path.normpath(os.path.join(base, '../data/Bruhin_et_al_2010.mat'))
matdata = loadmat(file_name)

df_full = pd.DataFrame()
for key in matdata.keys():
    if not '__' in key:
        df_full[key] = matdata[key].flatten()

## drop lotteries over losses
df = df_full[df_full['z2'] >= 0]

df_lottery = preprocess_df(df)
z1 = df_lottery['z1'].values
z2 = df_lottery['z2'].values
p1 = df_lottery['p1'].values
expval = df_lottery['expected_value'].values

base = os.path.dirname(os.path.abspath(__file__))
file_name = os.path.normpath(os.path.join(base, 'generated_ce.csv'))
df_ce = pd.read_csv(file_name)
ce_data = df_ce.values

# set search space for parameters
alpha_grid_size = 0.01
eta_grid_size = 0.01
alpha_lbd = 0.01
alpha_ubd = 1
eta_lbd = -1
eta_ubd = 5

alpha_grid = np.arange(alpha_lbd, alpha_ubd+alpha_grid_size, alpha_grid_size)
eta_grid = np.arange(eta_lbd, eta_ubd+eta_grid_size, eta_grid_size)
param_grid = cartesian_product(alpha_grid, eta_grid)

restrictiveness, stderr = compute_restrictiveness(z1, z2, p1, expval, ce_data, param_grid, closest=closest_DA)
print('restrictiveness: {}'.format(restrictiveness))
print('stderr: {}'.format(stderr))
